import sys
from math import sqrt, isnan
from math import log, pi, sqrt, copysign
from itertools import izip, groupby

from bessel import Bessel2NormFix

ALPHABET = 'actg'
PERIOD   = 3

def inv_xrange(n):
    for i in xrange(n):
        yield n - i - 1

def clog(f):
    if f > 0:
        return f * log(f)
    else:
        return 0.0

def mutual_inf_mx(m1, m2):
    I = 0
    for j in xrange(PERIOD):
        y1, y2 = 0, 0
        for c in ALPHABET:
            w1 = m1[j][c]
            w2 = m2[j][c]
            y1 += w1
            y2 += w2
            I += clog(w1) + clog(w2) - clog(w1 + w2)
        I += clog(y1 + y2) - clog(y1) - clog(y2)
    return I

def inf_mx(mx):
    I = 0
    y = dict((j,sum(mx[j][c] for c in ALPHABET)) for j in xrange(PERIOD))
    x = dict((c,sum(mx[j][c] for j in xrange(PERIOD))) for c in ALPHABET)
    L = float(sum(x.itervalues()))
    assert L == sum(y.itervalues())    
    for j in xrange(PERIOD):
        for c in ALPHABET:
            I += clog(mx[j][c])
    for j in xrange(PERIOD):
        I -= clog(y[j])
    for c in ALPHABET:
        I -= clog(x[c])
    I += clog(L)
    return I
    
def normalize_mx(mx):
    y = dict((j,sum(mx[j][c] for c in ALPHABET)) for j in xrange(PERIOD))
    x = dict((c,sum(mx[j][c] for j in xrange(PERIOD))) for c in ALPHABET)
    L = float(sum(x.itervalues()))
    assert L == sum(y.itervalues())    
    nmx = {}
    for j in xrange(PERIOD):
        nmx[j] = {}
        for c in ALPHABET:
            p = y[j] * x[c] / L / L
            nmx[j][c] = p and (mx[j][c] - p * L) / sqrt(L * p * (1 - p)) or 0.0
    return nmx

def norm_dist_mx(nmx1, nmx2):
    D = 0.0
    for j in xrange(PERIOD):
        for c in ALPHABET:
            D += Bessel2NormFix(nmx1[j][c] * nmx2[j][c])
    return D

def chi_dist_mx(nmx1, nmx2):
    D = 0.0
    for j in xrange(PERIOD):
        for c in ALPHABET:
            D += (nmx1[j][c] - nmx2[j][c])**2 / 2.0
    return D

def norm2prob(x, m=0.0, s=6.0):
    from scipy.stats import norm
    return norm.cdf(x, m, s)

def chi2prob(x, n_degrees=6):
    from scipy.stats import chi2    
    return chi2.cdf(x, n_degrees)    

def create_mx(seq):
    d = dict((j, dict((c,0) for c in ALPHABET)) for j in xrange(PERIOD))
    for i, c in enumerate(seq):
        if c in ALPHABET:
            d[i%PERIOD][c] += 1
    #print seq
    return d

def create_norm_mx(seq):
    return normalize_mx(create_mx(seq))

def create_base_seq(mx):
    return ''.join(
        ''.join(item) for item in izip(
            *[ ''.join(c*mx[j][c] for c in ALPHABET) for j in xrange(PERIOD) ]
        )
    )

def resize_mx(mx, length):
    L = sum(sum(mx[j][c] for c in ALPHABET) for j in xrange(PERIOD))
    multiplier = float(length) / L
    nmx = {}
    for j in xrange(PERIOD):
        nmx[j] = {}
        for c in ALPHABET:
            nmx[j][c] = int(round(mx[j][c] * multiplier))
    return nmx
    

def print_mx(mx):
    for c in ALPHABET:
        print "\t".join("%.3f" % mx[j][c] for j in xrange(PERIOD))


def print_dp_mx(dp_mx):
    TERM_WIDTH = 180
    COL_WIDTH  = 12
    COL_PER_LINE = TERM_WIDTH / COL_WIDTH
    N = len(dp_mx[0])
    L = len(dp_mx)
    for k,b in enumerate(xrange(0, L, COL_PER_LINE)):
        for i in xrange(N):
           col = " ".join("%7.2f %i" % dp_mx[j][i] for j in xrange(b, min(L,b+COL_PER_LINE)))
           print "%3i\t%s" % (k*COL_PER_LINE, col)
        print

def create_dp_mx(options, base_seq, seq):
    reward     = options.reward
    penalty    = options.penalty
    win_length = options.win_length
    length     = len(seq)
    seq        = seq.lower()
    
    base_mx = create_norm_mx(base_seq)
    
    bmxs = [dict(((i-j)%PERIOD, n) for i,n in base_mx.iteritems()) for j in xrange(PERIOD)]
    #print bmxs
    #for j in xrange(PERIOD):
    #    print '---------'
    #    print_mx(bmxs[j])
    
    dp_mx = [[(0,i) for i in xrange(PERIOD)]]
    #print_dp_mx(dp_mx)
    
    for x in xrange(0, length-win_length-1):
        #print x, "################################"
        mxs = [create_norm_mx(seq[x+j:x+j+win_length]) for j in xrange(PERIOD)]
        #print_mx(bmxs[x % PERIOD])
        #print '***'
        #print_mx(mxs[0])
        #print '***'
        deltas = [norm_dist_mx(bmxs[x%PERIOD], mxs[j]) for j in xrange(PERIOD)]
        #print 'deltas: ', ' | '.join(map(str, deltas))
        new_col = [
            max(
                (
                    (
                        max(dp_mx[x][j][0] + (i==j and deltas[j]-reward or -penalty), 0),
                        j
                    )
                    for j in (k%PERIOD for k in xrange(i, i+PERIOD))
                ),
                key = lambda x:x[0],
            )
            for i in xrange(PERIOD)
        ]
        dp_mx.append(new_col)
        #print "%5i: %s" % (x, ''.join("(%7.2f, %i), " % _i for _i in new_col))
        #print rewards
        #print_dp_mx(dp_mx)
    
    return dp_mx

def read_dp_mx(dp_mx):
    max_reward = float('-inf')
    for i in xrange(len(dp_mx)):
        for j in xrange(PERIOD):
            reward = dp_mx[i][j][0]
            if reward > max_reward:
                max_reward = reward
                max_i, max_j = i, j
    
    #print 'RES:', max_reward, max_i, max_j

    backtrace = []
    start_pos = 0
    if max_reward > 0:
        shift = max_j
        for i in inv_xrange(max_i+1):
            reward, shift = dp_mx[i][shift]
            #print i, reward, shift
            if reward <= 0:
                break
            backtrace.append(shift)
        backtrace.reverse()
        start_pos = max_i - len(backtrace) + 1
    return max_reward, start_pos, [(shift, len(list(grp))) for shift, grp in groupby(backtrace)]

def print_res(subseq, shifts):
    if shifts:
        raise Exception('Not implemented for shifts instead backtrace!')
        last_shift = backtrace[0]
        i = 0
        for j, shift in enumerate(backtrace):
            diff = (PERIOD - abs(shift - last_shift)) % PERIOD
            for k in xrange(diff):
                i += 1
                sys.stdout.write('-')
                if i % PERIOD == 0:
                    sys.stdout.write('|')
            sys.stdout.write(subseq[j])
            i += 1
            if i % PERIOD == 0:
                sys.stdout.write('|')
            last_shift = shift
    sys.stdout.write('\n')
